// --------------------------------------------------------------------------------------------------------------------
// <copyright file="UdpListener.cs" company="Microsoft Corporation">
// The MIT License (MIT)
// 
// Copyright (c) 2014, Microsoft Corporation
// 
// Permission is hereby granted, free of charge, to any person obtaining a copy
//  of this software and associated documentation files (the "Software"), to deal
//  in the Software without restriction, including without limitation the rights
//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//  copies of the Software, and to permit persons to whom the Software is
//  furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in
//  all copies or substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
//  THE SOFTWARE.
// </copyright>
// --------------------------------------------------------------------------------------------------------------------
namespace Microsoft.Robotics.Runtime
{
    using System;
    using System.Collections.Generic;
    using System.Diagnostics;
    using System.IO;
    using System.Net;
    using System.Net.Sockets;
    using System.Runtime.Serialization;
    using System.Threading;
    using System.Xml;
    using Microsoft.Robotics.IO;
    using Microsoft.Robotics.Tracing;

    /// <summary>
    /// A listener that receives a UDP message, deserializes it and sends to a forwarder.
    /// The class doesn't perform any chunking. Serialized Agent Messages must fit in 64K.
    /// </summary>
    internal class UdpListener : IDisposable
    {
        /// <summary>
        /// The weakly-typed forwarders that accept messages retrieved by this listener.
        /// </summary>
        private Dictionary<Type, IForwarder> forwarders;

        /// <summary>
        /// The set of sequenceIds we've seen last
        /// </summary>
        private Dictionary<Type, long> watermarks;

        /// <summary>
        /// The udp client used to send messages.
        /// </summary>
        private UdpClient updClient;

        /// <summary>
        /// Chunked UDP datagram receiver.
        /// </summary>
        private DataUnchunker unchunker;

        /// <summary>
        /// The data contract serializer to use.
        /// </summary>
        private DataContractSerializer serializer;

        /// <summary>
        /// The client address from which we accept messages.
        /// </summary>
        private IPAddress clientIp;

        /// <summary>
        /// The network port number to listen to.
        /// </summary>
        private int port;

        /// <summary>
        /// A simple lock.
        /// </summary>
        private SpinLock spinLock;

        /// <summary>
        /// The time when the parent pipeline element was created.
        /// </summary>
        private long pipelineCreationTime;

        /// <summary>
        /// Initializes a new instance of the UdpListener class.
        /// The instance starts listening right away to the specified port.
        /// </summary>
        /// <param name="port">local port to use when receiving messages.</param>
        /// <param name="clientIp">The client ip from which we accept messages.</param>
        /// <param name="forwarders">The set of forwarders to call when a message is received, indexed by type.</param>
        /// <param name="pipelineCreationTime">The time when the parent pipeline element was created.</param>
        public UdpListener(int port, IPAddress clientIp, Dictionary<Type, IForwarder> forwarders, long pipelineCreationTime)
        {
            this.forwarders = forwarders;
            this.pipelineCreationTime = pipelineCreationTime;
            this.watermarks = new Dictionary<Type, long>();
            foreach (Type messageType in forwarders.Keys)
            {
                this.watermarks[messageType] = 0;
            }

            this.clientIp = clientIp;
            this.port = port;
            this.serializer = new DataContractSerializer(
                typeof(AgentMessage), forwarders.Keys, Serializer.MaxItemsInTheObjectGraph, true, false, new SerializationSurrogate(false));

            this.unchunker = new DataUnchunker(UdpSender.MaxSerializedBytes, UdpSender.MaxUdpDatagramSize, _ => { }, _ => { });
        }

        /// <summary>
        /// Disposes of an instance of the <see cref="UdpListner"/> class.
        /// </summary>
        public void Dispose()
        {
            this.Dispose(true);
        }

        /// <summary>
        /// Starts listening for messages.
        /// </summary>
        public void Start()
        {
            if (this.updClient != null)
            {
                throw new InvalidOperationException("The listener is already started.");
            }

            this.updClient = new UdpClient(this.port);
            this.updClient.BeginReceive(this.ReceiveCallback, null);
        }

        /// <summary>
        /// Stops listening for messages.
        /// </summary>
        public void Stop()
        {
            if (this.updClient == null)
            {
                throw new InvalidOperationException("The listener is already stopped.");
            }

            // lock access to the inner UdpClient
            bool lockTaken = false;
            this.spinLock.Enter(ref lockTaken);

            if (!lockTaken)
            {
                throw new InvalidOperationException("Could not aquire the lock.");
            }

            // we got the lock
            try
            {
                this.updClient.Close();
                this.updClient = null;
            }
            finally
            {
                this.spinLock.Exit();
            }
        }

        /// <summary>
        /// Async callback called when a new message is available
        /// </summary>
        /// <param name="ar">The async result</param>
        private void ReceiveCallback(IAsyncResult ar)
        {
            ////Console.WriteLine("Receive: {0}", this.port);
            IPEndPoint clientEndpoint = new IPEndPoint(0, 0);

            // lock access to the inner UdpClient
            // The purpose of the spin lock is to block the shutdown code path (Stop()) while we are using the socket.
            // We use a SpinLock rather than Monitor because SpinLock is much cheaper and at steady state this thread is the only thread that needs the lock at any time.
            // The ReceiveCallback method is not re-entrant on different threads, so the only code path that might wait for us is Stop(), 
            // and we really don't care if the thread calling Stop will spin waiting for a long time, since Stop is called only at shutdown time (i.e. very rarely).
            // One special case: udpClient.BeginReceive is sometimes not async (presumably when a new packet is available), resulting in a sync call to ReceiveCallback.
            // To cover this case, we simply verify if the lock is already taken by the current thread.
            bool lockTaken = false;
            if (!this.spinLock.IsHeldByCurrentThread)
            {
                this.spinLock.Enter(ref lockTaken);
                if (!lockTaken)
                {
                    throw new InvalidOperationException("Could not aquire the lock.");
                }
            }

            // got the lock, now get the bytes off the wire and turn them into a message
            try
            {
                // are we shutting down?
                if (this.updClient == null)
                {
                    return;
                }

                if (this.unchunker.Receive(this.updClient.EndReceive(ar, ref clientEndpoint)))
                {
                    // if this message is not from the IP we expect, ignore it
                    if (clientEndpoint.Address.Equals(this.clientIp))
                    {
                        // deserialize the agent message
                        AgentMessage result = null;

                        using (
                            XmlDictionaryReader reader = XmlDictionaryReader.CreateBinaryReader(
                                this.unchunker.Payload, 0, this.unchunker.Length, XmlDictionaryReaderQuotas.Max))
                        {
                            result = (AgentMessage)this.serializer.ReadObject(reader);

                            // replace the timestamp with our time to make it compatible with local messages
                            result.Timestamp = AgentMessage.GetCurrentTime() - this.pipelineCreationTime;
                        }

                        // if this message is too old, ignore it
                        Type resultType = result.GetType();
                        if (result.SequenceId > this.watermarks[resultType])
                        {
                            // pass the deserialized object to the consumer
                            IForwarder forwarder = this.forwarders[resultType];
                            forwarder.Post(result);

                            // bump up the watermark
                            this.watermarks[resultType] = result.SequenceId;
                        }
                    }
                }
            }
            catch (Exception e)
            {
                TraceOut.Warning(TraceContexts.Common, "Exception in Listener: {0}", e.Message);

                // Bug 2089 - remove this debug code once resolved
                ////byte[] buf = new byte[this.unchunker.Payload.Length];
                ////Array.Copy(this.unchunker.Payload, buf, buf.Length);
                ////File.WriteAllBytes(@"c:\temp\payload.txt", buf);
            }
            finally
            {
                // keep listening
                if (null != this.updClient)
                {
                    this.updClient.BeginReceive(this.ReceiveCallback, null);
                }

                if (lockTaken)
                {
                    this.spinLock.Exit();
                }
            }
        }

        /// <summary>
        /// Disposes of an instance of the <see cref="UdpListner"/> class.
        /// </summary>
        /// <param name="disposing">True if disposing, false if finalizing.</param>
        private void Dispose(bool disposing)
        {
            if (disposing)
            {
                if (this.updClient != null)
                {
                    this.Stop();
                }

                GC.SuppressFinalize(this);
            }
        }
    }
}
